/**
 * Copyright (c) 2023 murenchao
 * taomu is licensed under Mulan PubL v2.
 * You can use this software according to the terms and conditions of the Mulan PubL v2.
 * You may obtain a copy of Mulan PubL v2 at:
 *       http://license.coscl.org.cn/MulanPubL-2.0
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PubL v2 for more details.
 */
package cool.taomu.utils.crypto.util

import cool.taomu.utils.CreateTempFile
import cool.taomu.utils.crypto.Base64
import java.io.ByteArrayInputStream
import java.io.ByteArrayOutputStream
import java.io.File
import java.nio.charset.Charset
import java.security.SecureRandom
import java.security.Security
import org.apache.commons.io.FileUtils
import org.bouncycastle.jce.provider.BouncyCastleProvider
import org.bouncycastle.openpgp.PGPCompressedData
import org.bouncycastle.openpgp.PGPCompressedDataGenerator
import org.bouncycastle.openpgp.PGPEncryptedData
import org.bouncycastle.openpgp.PGPEncryptedDataGenerator
import org.bouncycastle.openpgp.PGPEncryptedDataList
import org.bouncycastle.openpgp.PGPLiteralData
import org.bouncycastle.openpgp.PGPOnePassSignatureList
import org.bouncycastle.openpgp.PGPPublicKeyEncryptedData
import org.bouncycastle.openpgp.PGPPublicKeyRingCollection
import org.bouncycastle.openpgp.PGPSecretKeyRingCollection
import org.bouncycastle.openpgp.PGPSignatureList
import org.bouncycastle.openpgp.PGPUtil
import org.bouncycastle.openpgp.jcajce.JcaPGPObjectFactory
import org.bouncycastle.openpgp.operator.jcajce.JcaKeyFingerprintCalculator
import org.bouncycastle.openpgp.operator.jcajce.JcePBESecretKeyDecryptorBuilder
import org.bouncycastle.openpgp.operator.jcajce.JcePGPDataEncryptorBuilder
import org.bouncycastle.openpgp.operator.jcajce.JcePublicKeyDataDecryptorFactoryBuilder
import org.bouncycastle.openpgp.operator.jcajce.JcePublicKeyKeyEncryptionMethodGenerator
import org.bouncycastle.util.io.Streams
import org.slf4j.LoggerFactory

class PgpUtils {
	val static LOG = LoggerFactory.getLogger(PgpUtils);

	def static void main(String[] args) {
		Security.addProvider(new BouncyCastleProvider());
		var String pk = FileUtils.readFileToString(new File("src/test/resources/publicKey1.txt"),
			Charset.forName("UTF-8"));
		var data = PgpUtils.encrypt("Hello World!!!! ".bytes, pk, true);
		println(new String(data, "UTF-8"));
		var String prk = FileUtils.readFileToString(new File("src/test/resources/privateKey1.txt"),
			Charset.forName("UTF-8"));
		var d = PgpUtils.decrypt(data, prk, "123456789");
		println("result : " + new String(d, "UTF-8"));
	}

	def static decrypt(byte[] data, String privateKey, String password) {

		try(var bais = new ByteArrayInputStream(data)) {
			var input = PGPUtil.getDecoderStream(bais);
			var jpof = new JcaPGPObjectFactory(input);
			var o = jpof.nextObject;
			var PGPEncryptedDataList enc;
			if (o instanceof PGPEncryptedDataList) {
				enc = o;
			} else {
				enc = jpof.nextObject as PGPEncryptedDataList;
			}
			// 获取privatekey
			var pkeyBase64 = new Base64(privateKey.bytes).decode();
			val pgpSec = new PGPSecretKeyRingCollection(pkeyBase64, new JcaKeyFingerprintCalculator());
			var edo = enc.encryptedDataObjects
			var pbe = edo.findFirst [ k |
				pgpSec.getSecretKey((k as PGPPublicKeyEncryptedData).keyID) !== null
			] as PGPPublicKeyEncryptedData;
			var secKey = pgpSec.getSecretKey(pbe.keyID);
			var jpekdb = new JcePBESecretKeyDecryptorBuilder().setProvider("BC").build(password.toCharArray);
			var sKey = secKey.extractPrivateKey(jpekdb);
			var clear = pbe.getDataStream(new JcePublicKeyDataDecryptorFactoryBuilder().setProvider("BC").build(sKey));
			var plain = new JcaPGPObjectFactory(clear);
			var message = plain.nextObject
			if (input !== null) {
				input.close();
			}
			try(var out = new ByteArrayOutputStream()) {
				if (message instanceof PGPCompressedData) {
					var fact = new JcaPGPObjectFactory(message.dataStream);
					message = fact.nextObject;
				}
				if (message instanceof PGPLiteralData) {
					Streams.pipeAll(message.inputStream, out);
				} else if (message instanceof PGPOnePassSignatureList) {
					LOG.info("encrypted message contains a signed message - not literal data.")
				} else if (message instanceof PGPSignatureList) {
					LOG.info("encrypted message contains a signed message - not literal data.");
				}
				if (pbe.isIntegrityProtected) {
					if (!pbe.verify) {
						LOG.info("message failed integrity check");
					} else {
						LOG.info("message integrity check passed");
					}
				} else {
					LOG.info("no message integrity check");
				}
				return out.toByteArray
			}
		}
	}

	def static encrypt(byte[] data, String publicKey, boolean withIntegrityCheck) {
		// 获取publiekey
		var bais = new ByteArrayInputStream(new Base64(publicKey.bytes).decode);
		var input = PGPUtil.getDecoderStream(bais);
		var pkrc = new PGPPublicKeyRingCollection(input, new JcaKeyFingerprintCalculator);
		var encKey = pkrc.keyRings.findFirst [
			it.publicKeys.findFirst[it.isEncryptionKey].isEncryptionKey;
		].publicKey;
		var file = CreateTempFile.create(data, "UTF-8");
		try(var baos = new ByteArrayOutputStream()) {
			var cdg = new PGPCompressedDataGenerator(PGPCompressedData.ZIP);
			PGPUtil.writeFileToLiteralData(cdg.open(baos), 'b', file);
			cdg.close;
			var bytes = baos.toByteArray;
			var jpdeb = new JcePGPDataEncryptorBuilder(PGPEncryptedData.CAST5);
			jpdeb = jpdeb.setWithIntegrityPacket(withIntegrityCheck);
			var bc = jpdeb.setSecureRandom(new SecureRandom()).setProvider("BC");
			var encGen = new PGPEncryptedDataGenerator(bc);
			encGen.addMethod(new JcePublicKeyKeyEncryptionMethodGenerator(encKey).provider = "BC");
			try(var out = new ByteArrayOutputStream()) {
				var cout = encGen.open(out, bytes.length);
				cout.write(bytes);
				cout.close();
				return out.toByteArray;
			}

		} finally {
			file.delete
		}
	}
}
